import sys
import os
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from io import BytesIO
import cv2
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QGridLayout, QLabel, QHBoxLayout, QVBoxLayout
)
from PyQt5.QtGui import QPixmap, QImage, QPen, QPainter, QMouseEvent, QKeyEvent, QWheelEvent, QColor, QTransform, QFontMetrics, QDragEnterEvent, QDropEvent
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtCore import Qt, QPoint, QPointF, QRect, QRectF, QSize, QSizeF, pyqtSignal, QUrl
from PyQt5.QtCore import QTimer, QEasingCurve, QPropertyAnimation, QRect, QPoint
from PyQt5.QtWidgets import QLabel, QDialog, QDialogButtonBox, QTextEdit
from PyQt5.QtGui import QFontDatabase
from enum import Enum
import image_metric


class UsageDialog(QDialog):
    def __init__(self, usage_text: str, parent=None):
        super().__init__(parent)
        self.setWindowTitle("使用说明")
        self.resize(720, 520)

        layout = QVBoxLayout(self)

        text = QTextEdit(self)
        text.setReadOnly(True)
        mono = QFontDatabase.systemFont(QFontDatabase.FixedFont)
        text.setFont(mono)
        text.setPlainText(usage_text)
        text.setLineWrapMode(QTextEdit.WidgetWidth)
        layout.addWidget(text)

        buttons = QDialogButtonBox(QDialogButtonBox.Close, parent=self)
        buttons.rejected.connect(self.reject)
        buttons.accepted.connect(self.accept)
        layout.addWidget(buttons)
class ToastWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowFlags(
            Qt.ToolTip | Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint
        )
        self.setAttribute(Qt.WA_TranslucentBackground, True)

        self._label = QLabel(self)
        self._label.setStyleSheet("""
            QLabel {
                color: white;
                background-color: rgba(0, 0, 0, 180);
                border-radius: 6px;
                padding: 8px 12px;
                font-size: 12pt;
            }
        """)
        self._label.setAlignment(Qt.AlignCenter)

        self._timer = QTimer(self)
        self._timer.setSingleShot(True)
        self._timer.timeout.connect(self._on_timeout)

        self._fade_anim = QPropertyAnimation(self, b"windowOpacity", self)
        self._fade_anim.setEasingCurve(QEasingCurve.InOutQuad)

        self._duration_ms = 500

    def _safe_disconnect_finished(self):
        try:
            self._fade_anim.finished.disconnect(self.hide)
        except Exception:
            pass
        try:
            self._fade_anim.finished.disconnect(self._on_fade_out_finished)
        except Exception:
            pass

    def show_message(self, message: str, duration_ms: int = 500):
        # 1) 清理上一次的状态：停计时器、停动画、断开 finished 旧连接
        self._timer.stop()
        self._fade_anim.stop()
        self._safe_disconnect_finished()

        self._duration_ms = duration_ms
        self._label.setText(message)
        self._label.adjustSize()

        margin = 12
        w = self._label.width() + margin
        h = self._label.height() + margin
        self.resize(w, h)
        self._label.move((w - self._label.width()) // 2,
                         (h - self._label.height()) // 2)

        # 2) 计算位置（使用全局坐标，避免窗口装饰误差）
        if self.parent() is not None:
            pw = self.parent().width()
            ph = self.parent().height()
            # 父窗口客户区中心的屏幕坐标
            local_pos = QPoint((pw - w) // 2, int(ph * 0.85) - h // 2)
            global_pos = self.parent().mapToGlobal(local_pos)
            x, y = global_pos.x(), global_pos.y()
        else:
            screen = QApplication.primaryScreen().availableGeometry()
            x = screen.x() + (screen.width() - w) // 2
            y = screen.y() + int(screen.height() * 0.85) - h // 2

        self.move(x, y)

        # 3) 渐入
        self.setWindowOpacity(0.0)
        self.show()
        self._fade_anim.stop()
        self._fade_anim.setDuration(120)
        self._fade_anim.setStartValue(0.0)
        self._fade_anim.setEndValue(1.0)
        self._fade_anim.start()

        # 4) 倒计时后渐出
        self._timer.start(self._duration_ms)

    def _on_timeout(self):
        # 渐出前先清理并仅连接一次结束回调
        self._fade_anim.stop()
        self._safe_disconnect_finished()
        self._fade_anim.setDuration(180)
        self._fade_anim.setStartValue(1.0)
        self._fade_anim.setEndValue(0.0)
        self._fade_anim.finished.connect(self._on_fade_out_finished)
        self._fade_anim.start()

    def _on_fade_out_finished(self):
        # 结束后隐藏，并断开连接，防止下次叠加
        self._safe_disconnect_finished()
        self.hide()

class ImageViewer(QWidget):
    offset_scale_changed = pyqtSignal(QPointF, float)
    position_changed = pyqtSignal(QPointF)
    rectangle_changed = pyqtSignal(QRectF)

    def __init__(self, image_path, name):
        super().__init__()
        self.image_path = image_path
        self.name = name
        self.is_show_normalization = False
        self.color_space = 'yuv'
        self.rotation_angle = 0
        self.initUI()

    def initUI(self):
        self.scale = 1.0
        self.offset = QPointF(0, 0)
        self.min_scale = 0.1
        self.max_scale = 40.0
        self.press_start_pos = QPointF(0, 0)
        self.move_pos = QPointF(0, 0)
        self.select_pos = None
        self.select_rect = None

        self.original_image = QImage(self.image_path)
        if self.original_image.isNull():
            print(f"Error: Unable to load image {self.image_path}.")
            sys.exit(1)

        self.image = self.original_image.copy()
        self.pixmap = QPixmap.fromImage(self.image)
        self.play_pixmap = self.pixmap.copy()

        self.info_bar = QWidget(self)
        self.info_bar.setGeometry(0, 0, self.width(), 30)
        self.info_bar.setMaximumHeight(30)
        self.info_bar.setStyleSheet("background-color: rgba(0, 0, 0, 100);")
        info_layout = QHBoxLayout()
        self.name_label = QLabel(self.name)
        self.name_label.setStyleSheet("color: yellow; background-color: transparent;")
        self.name_label.setToolTip(self.image_path)
        self.setStyleSheet("""QToolTip {
        color: #ffeb3b;
        background-color: rgba(0,0,0,200);
        border: 1px solid #888;
        padding: 4px;
        font-size: 12px;
        font-family: "Microsoft YaHei", "Segoe UI", Arial;
        }""")

        self.position_color_label = QLabel("Position: (0, 0), YUV:(0,0,0)")
        self.position_color_label.setStyleSheet("color: yellow; background-color: transparent;")
        info_layout.addWidget(self.name_label)
        info_layout.addStretch()
        info_layout.addWidget(self.position_color_label)
        self.info_bar.setLayout(info_layout)

        self._set_name_label_text()

        self.hist_label = QLabel(self)
        self.hist_label.setStyleSheet("background-color: rgba(255, 255, 255, 200);")
        self.hist_label.setAlignment(Qt.AlignCenter)
        self.hist_label.hide()

        self.setMinimumSize(200, 200)

    def _shorten_path_for_label(self, full_path, keep_levels=2):
        dir_path = os.path.dirname(full_path)
        norm = dir_path.replace('\\', '/')
        parts = [p for p in norm.split('/') if p]
        if len(parts) == 0:
            return ""
        if len(parts) <= keep_levels:
            short = dir_path
        else:
            short = '.../' + '/'.join(parts[-keep_levels:])
        return short

    def _set_name_label_text(self, max_width_px=None):
        base = os.path.basename(self.image_path) if self.image_path else self.name
        short_parent = self._shorten_path_for_label(self.image_path, keep_levels=2)
        show_text = base if not short_parent else f"{base}  ({short_parent})"
        label_width = max_width_px if max_width_px is not None else max(120, int(self.width() * 0.45))
        fm = QFontMetrics(self.name_label.font())
        elided = fm.elidedText(show_text, Qt.ElideMiddle, label_width)
        self.name_label.setText(elided)
        self.name_label.setToolTip(self.image_path)

    def rotateImage(self, clockwise=True):
        if clockwise:
            self.rotation_angle = (self.rotation_angle + 90) % 360
        else:
            self.rotation_angle = (self.rotation_angle - 90) % 360

        transform = QTransform()
        transform.rotate(self.rotation_angle)
        self.image = self.original_image.transformed(transform, Qt.SmoothTransformation)
        self.pixmap = QPixmap.fromImage(self.image)
        self.play_pixmap = self.pixmap.copy()

        self.select_pos = None
        self.select_rect = None
        self.hist_label.hide()
        self.position_color_label.setText("Position: (0, 0), YUV:(0,0,0)")
        self.offset = QPointF(0, 0)
        self.scale = 1.0
        self.update()

    def resizeEvent(self, event):
        self.info_bar.setGeometry(0, 0, self.width(), 40)
        hist_width = 300
        hist_height = 200
        self.hist_label.setGeometry(
            self.width() - hist_width - 10,
            self.height() - hist_height - 10,
            hist_width,
            hist_height
        )
        available = max(120, int(self.width() * 0.45))
        self._set_name_label_text(max_width_px=available)
        super().resizeEvent(event)

    def paintEvent(self, event):
        painter = QPainter(self)
        scaled_size = QSizeF(self.pixmap.size()) * self.scale
        img_rect = QRectF(self.offset, scaled_size)
        visible_rect = QRectF(0, 0, self.width(), self.height())
        draw_rect = img_rect.intersected(visible_rect)
        if not draw_rect.isEmpty():
            source_rect = QRectF(
                (draw_rect.x() - self.offset.x()) / self.scale,
                (draw_rect.y() - self.offset.y()) / self.scale,
                draw_rect.width() / self.scale,
                draw_rect.height() / self.scale
            )
            painter.drawPixmap(draw_rect, self.play_pixmap, source_rect)

    def mouseReleaseEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton:
            pos = event.pos()
            delta = pos - self.press_start_pos
            dis = abs(delta.x()) + abs(delta.y())
            if dis < 3:
                img_x = (self.press_start_pos.x() - self.offset.x()) / self.scale
                img_y = (self.press_start_pos.y() - self.offset.y()) / self.scale
                if 0 <= img_x < self.image.width() and 0 <= img_y < self.image.height():
                    self.position_changed.emit(QPointF(img_x, img_y))
            else:
                img_start_x = (self.press_start_pos.x() - self.offset.x()) / self.scale
                img_start_y = (self.press_start_pos.y() - self.offset.y()) / self.scale
                img_end_x = (pos.x() - self.offset.x()) / self.scale
                img_end_y = (pos.y() - self.offset.y()) / self.scale
                if 0 <= img_start_x < self.image.width() and 0 <= img_start_y < self.image.height() and 0 <= img_end_x < self.image.width() and 0 <= img_end_y < self.image.height():
                    left = min(img_start_x, img_end_x)
                    top = min(img_start_y, img_end_y)
                    right = max(img_start_x, img_end_x)
                    bottom = max(img_start_y, img_end_y)
                    rect = QRectF(QPointF(left, top), QPointF(right, bottom))
                    self.rectangle_changed.emit(rect)

    def mousePressEvent(self, event: QMouseEvent):
        self.press_start_pos = event.pos()
        self.move_pos = event.pos()

    def mouseMoveEvent(self, event: QMouseEvent):
        if event.buttons() & Qt.RightButton:
            delta = event.pos() - self.move_pos
            self.offset += delta
            self.move_pos = event.pos()
            self.update()
            self.offset_scale_changed.emit(self.offset, self.scale)
        elif event.buttons() & Qt.LeftButton:
            pos = event.pos()
            delta = pos - self.press_start_pos
            dis = abs(delta.x()) + abs(delta.y())
            if dis > 3:
                img_start_x = (self.press_start_pos.x() - self.offset.x()) / self.scale
                img_start_y = (self.press_start_pos.y() - self.offset.y()) / self.scale
                img_end_x = (pos.x() - self.offset.x()) / self.scale
                img_end_y = (pos.y() - self.offset.y()) / self.scale
                rect = QRectF(img_start_x, img_start_y, img_end_x - img_start_x, img_end_y - img_start_y)
                self.updateRectangle(rect)

    def wheelEvent(self, event: QWheelEvent):
        delta = event.angleDelta().y()
        old_center = QPointF(
            (self.width() / 2.0 - self.offset.x()) / self.scale,
            (self.height() / 2.0 - self.offset.y()) / self.scale
        )
        if delta > 0:
            self.scale = min(self.scale * 1.1, self.max_scale)
        else:
            self.scale = max(self.scale * 0.9, self.min_scale)
        self.offset = QPointF(
            self.width() / 2.0 - old_center.x() * self.scale,
            self.height() / 2.0 - old_center.y() * self.scale
        )
        self.offset_scale_changed.emit(self.offset, self.scale)
        self.update()

    def updateGeometry(self, offset: QPointF, scale: float):
        self.offset = offset
        self.scale = scale
        self.update()

    def rgbToHsv(self, rgb):
        rgb = np.asarray(rgb, dtype=float).reshape(-1)
        if rgb.size != 3:
            raise ValueError("rgb must be an iterable of length 3, e.g. (R, G, B).")
        r, g, b = rgb
        if r > 1.0 or g > 1.0 or b > 1.0:
            r, g, b = r / 255.0, g / 255.0, b / 255.0
        cmax = max(r, g, b)
        cmin = min(r, g, b)
        delta = cmax - cmin
        eps = 1e-10
        if delta < eps:
            h = 0.0
        elif cmax == r:
            h = 60.0 * (((g - b) / (delta + eps)) % 6.0)
        elif cmax == g:
            h = 60.0 * (((b - r) / (delta + eps)) + 2.0)
        else:
            h = 60.0 * (((r - g) / (delta + eps)) + 4.0)
        if cmax < eps:
            s = 0.0
        else:
            s = delta / (cmax + eps)
        v = cmax
        h = h % 360.0
        return np.array([h, s, v], dtype=float)

    def rgbToYuv(self, rgb):
        rgb = rgb / 255.0
        transform_matrix = np.array([
            [0.299, 0.587, 0.114],
            [-0.14713, -0.28886, 0.436],
            [0.615, -0.51499, -0.10001]
        ])
        yuv = np.dot(transform_matrix, rgb)
        yuv[1] += 0.5
        yuv[2] += 0.5
        yuv = np.clip(yuv * 255, 0, 255).astype(np.uint8)
        return yuv

    def updatePosition(self, pos: QPointF):
        self.select_pos = pos
        self.select_rect = None
        self.hist_label.hide()
        x, y = int(pos.x()), int(pos.y())
        if 0 <= x < self.image.width() and 0 <= y < self.image.height():
            color = self.image.pixelColor(x, y)
            rgba = np.array([color.red(), color.green(), color.blue(), color.alpha()])
            if self.color_space == 'yuv':
                yuv = self.rgbToYuv(rgba[:3])
                if self.is_show_normalization:
                    yuv = yuv / 255.0
                    self.position_color_label.setText(
                        f"Position: ({x}, {y}), YUV:({yuv[0]:0.2f}, {yuv[1]:0.2f}, {yuv[2]:0.2f}), "
                    )
                else:
                    self.position_color_label.setText(
                        f"Position: ({x}, {y}), YUV:({int(yuv[0])}, {int(yuv[1])}, {int(yuv[2])}), "
                    )
            elif self.color_space == 'rgba':
                if self.is_show_normalization:
                    rgba = rgba / 255.0
                    self.position_color_label.setText(
                        f"Position: ({x}, {y}), RGBA:({rgba[0]:0.2f}, {rgba[1]:0.2f}, {rgba[2]:0.2f}, {rgba[3]:0.2f}), "
                    )
                else:
                    self.position_color_label.setText(
                        f"Position: ({x}, {y}), RGBA:({rgba[0]}, {rgba[1]}, {rgba[2]}, {rgba[3]}), "
                    )
            else:
                hsv_val = self.rgbToHsv(rgba[:3].astype(np.float32))
                self.position_color_label.setText(
                    f"Position: ({x}, {y}), HSV:({hsv_val[0]:.1f}, {hsv_val[1]:.3f}, {hsv_val[2]:.3f}), "
                )
        else:
            self.position_color_label.setText("Position: (0, 0), RGBA: (0, 0, 0, 0)")

        self.play_pixmap = self.pixmap.copy()
        line_lenght = 10
        line_width = 1
        painter = QPainter(self.play_pixmap)
        pen = QPen(QColor(255, 0, 0, 100), line_width)
        painter.setPen(pen)
        painter.drawLine(x - line_lenght, y, x - 1, y)
        painter.drawLine(x + 1, y, x + line_lenght, y)
        painter.drawLine(x, y - line_lenght, x, y - 1)
        painter.drawLine(x, y + 1, x, y + line_lenght)
        painter.end()
        self.update()

    def calYUVRange(self, image):
        yuv = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
        y_range = (np.min(yuv[:, :, 0]), np.max(yuv[:, :, 0]))
        u_range = (np.min(yuv[:, :, 1]), np.max(yuv[:, :, 1]))
        v_range = (np.min(yuv[:, :, 2]), np.max(yuv[:, :, 2]))
        return y_range, u_range, v_range

    def calMultiChannelHist(self, image, channels, channel_names, output_width, output_height, dpi=100):
        fig_width = output_width / dpi
        fig_height = output_height / dpi
        plt.figure(figsize=(fig_width, fig_height), dpi=dpi)
        for i, channel in enumerate(channels):
            plt.subplot(1, len(channels), i + 1)
            channel_data = image[:, :, channel]
            channel_mean = float(np.mean(channel_data))
            data_min = channel_data.min()
            data_max = channel_data.max()
            if data_min == data_max:
                data_min = data_min - 0.5
                data_max = data_max + 0.5
            hist, bins = np.histogram(channel_data, bins=20, density=True, range=(data_min, data_max))
            plt.bar(bins[:-1], hist, width=np.diff(bins), align="edge", color='skyblue', edgecolor='black')
            ticks = np.linspace(data_min, data_max, num=5, dtype=np.float32)
            plt.xticks(ticks, rotation=-45)
            if np.issubdtype(channel_data.dtype, np.integer):
                plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%d'))
            else:
                if data_max - data_min < 0.5:
                    plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%.2f'))
                else:
                    plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
            plt.title(f"{channel_names[i]}({channel_mean:.1f})")
            plt.grid(axis='y', alpha=0.5)

        plt.tight_layout()
        buf = BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)
        buf.seek(0)
        hist_img = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        hist_img = cv2.imdecode(hist_img, cv2.IMREAD_COLOR)
        plt.close()
        hist_img = cv2.resize(hist_img, (output_width, output_height))
        return hist_img

    def updateRectangle(self, rect: QRectF, statistics=False):
        self.select_rect = rect
        self.select_pos = None

        if statistics:
            qimage = self.image.copy(rect.toRect()).convertToFormat(QImage.Format_ARGB32)
            if not qimage.isNull():
                ptr = qimage.bits()
                ptr.setsize(qimage.byteCount())
                arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
                if self.color_space == 'yuv':
                    arr_rgb = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGB)
                    yuv = cv2.cvtColor(arr_rgb, cv2.COLOR_RGB2YUV)
                    if self.is_show_normalization:
                        yuv = yuv.astype(np.float32) / 255.0
                    y_range = (np.min(yuv[:, :, 0]), np.max(yuv[:, :, 0]))
                    u_range = (np.min(yuv[:, :, 1]), np.max(yuv[:, :, 1]))
                    v_range = (np.min(yuv[:, :, 2]), np.max(yuv[:, :, 2]))
                    y_std = np.std(yuv[:, :, 0])
                    u_std = np.std(yuv[:, :, 1])
                    v_std = np.std(yuv[:, :, 2])
                    if self.is_show_normalization:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"yuv: ({y_range[0]:.2f}-{y_range[1]:.2f}, {u_range[0]:.2f}-{u_range[1]:.2f}, {v_range[0]:.2f}-{v_range[1]:.2f}), "
                            f"std: ({y_std:.2f}, {u_std:.2f}, {v_std:.2f})"
                        )
                    else:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"yuv: ({y_range[0]}-{y_range[1]}, {u_range[0]}-{u_range[1]}, {v_range[0]}-{v_range[1]}), "
                            f"std: ({y_std:.2f}, {u_std:.2f}, {v_std:.2f})"
                        )
                elif self.color_space == 'rgba':
                    rgba = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGBA)
                    if self.is_show_normalization:
                        rgba = rgba.astype(np.float32) / 255.0
                    r_range = (np.min(rgba[:, :, 0]), np.max(rgba[:, :, 0]))
                    g_range = (np.min(rgba[:, :, 1]), np.max(rgba[:, :, 1]))
                    b_range = (np.min(rgba[:, :, 2]), np.max(rgba[:, :, 2]))
                    a_range = (np.min(rgba[:, :, 3]), np.max(rgba[:, :, 3]))
                    r_std = np.std(rgba[:, :, 0])
                    g_std = np.std(rgba[:, :, 1])
                    b_std = np.std(rgba[:, :, 2])
                    a_std = np.std(rgba[:, :, 3])
                    if self.is_show_normalization:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"rgba: ({r_range[0]:.2f}-{r_range[1]:.2f}, {g_range[0]:.2f}-{g_range[1]:.2f}, {b_range[0]:.2f}-{b_range[1]:.2f}, {a_range[0]:.2f}-{a_range[1]:.2f}), "
                            f"std: ({r_std:.2f}, {g_std:.2f}, {b_std:.2f}, {a_std:.2f})"
                        )
                    else:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"rgba: ({r_range[0]}-{r_range[1]}, {g_range[0]}-{g_range[1]}, {b_range[0]}-{b_range[1]}, {a_range[0]}-{a_range[1]}), "
                            f"std: ({r_std:.2f}, {g_std:.2f}, {b_std:.2f}, {a_std:.2f})"
                        )
                else:
                    rgb = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGB).astype(np.float32) / 255.0
                    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
                    h_range = (np.min(hsv[:, :, 0]), np.max(hsv[:, :, 0]))
                    s_range = (np.min(hsv[:, :, 1]), np.max(hsv[:, :, 1]))
                    v_range = (np.min(hsv[:, :, 2]), np.max(hsv[:, :, 2]))
                    h_std = np.std(hsv[:, :, 0])
                    s_std = np.std(hsv[:, :, 1])
                    v_std = np.std(hsv[:, :, 2])
                    info_text = (
                        f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                        f"hsv: ({h_range[0]:.1f}-{h_range[1]:.1f}, {s_range[0]:.2f}-{s_range[1]:.2f}, {v_range[0]:.2f}-{v_range[1]:.2f}), "
                        f"std: ({h_std:.2f}, {s_std:.2f}, {v_std:.2f})"
                    )
                self.position_color_label.setText(info_text)

        self.play_pixmap = self.pixmap.copy()
        line_width = 1
        painter = QPainter(self.play_pixmap)
        pen = QPen(QColor(255, 0, 0, 100), line_width)
        painter.setPen(pen)
        painter.drawRect(rect)
        painter.end()
        self.update()

    def updateHist(self, rect: QRectF):
        qimage = self.image.copy(rect.toRect()).convertToFormat(QImage.Format_ARGB32)
        if not qimage.isNull():
            ptr = qimage.bits()
            ptr.setsize(qimage.byteCount())
            arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
            rgba = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGBA)
            if self.color_space == 'yuv':
                yuv = cv2.cvtColor(rgba[:, :, :3], cv2.COLOR_RGB2YUV)
                if self.is_show_normalization:
                    yuv = yuv.astype(np.float32) / 255.0
                hist_img = self.calMultiChannelHist(
                    yuv, [0, 1, 2], ['Y', 'U', 'V'], output_width=600, output_height=200
                )
            elif self.color_space == 'rgba':
                data = rgba
                if self.is_show_normalization:
                    data = data.astype(np.float32) / 255.0
                hist_img = self.calMultiChannelHist(
                    data[:, :, :3], [0, 1, 2], ['R', 'G', 'B'], output_width=600, output_height=200
                )
            else:
                rgbf = rgba[:, :, :3].astype(np.float32) / 255.0
                hsv = cv2.cvtColor(rgbf, cv2.COLOR_RGB2HSV)
                hist_img = self.calMultiChannelHist(
                    hsv, [0, 1, 2], ['H', 'S', 'V'], output_width=600, output_height=200
                )

            height, width, channel = hist_img.shape
            bytes_per_line = 3 * width
            q_hist_img = QImage(hist_img.data, width, height, bytes_per_line, QImage.Format_RGB888)
            self.hist_label.setPixmap(QPixmap.fromImage(q_hist_img))
            self.hist_label.setGeometry(self.width() - width - 5, self.height() - height - 5, width, height)
            self.hist_label.show()

    def updateMTF(self, rect: QRectF):
        qimage = self.image.copy(rect.toRect()).convertToFormat(QImage.Format_ARGB32)
        if not qimage.isNull():
            ptr = qimage.bits()
            ptr.setsize(qimage.byteCount())
            arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
            rgb = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGB)
            yuv = cv2.cvtColor(rgb, cv2.COLOR_RGB2YUV)
            luma = yuv[:, :, 0]
            mtf_50, nyquist = image_metric.compute_mtf(luma)
            if mtf_50 >= 0 and nyquist >= 0:
                self.position_color_label.setText(f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), MTF50:{mtf_50:.2f}, Nyquist:{nyquist:.2f}%")
            else:
                self.position_color_label.setText(f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), MTF计算失败")

    def updatePSNR(self, rect: QRectF):
        qimage = self.image.copy(rect.toRect()).convertToFormat(QImage.Format_ARGB32)
        if not qimage.isNull():
            ptr = qimage.bits()
            ptr.setsize(qimage.byteCount())
            arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
            rgba = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGBA)
            if self.color_space == 'yuv':
                yuv = cv2.cvtColor(rgba[:, :, :3], cv2.COLOR_RGB2YUV)
                if self.is_show_normalization:
                    yuv = yuv.astype(np.float32) / 255.0
                src_img = yuv
            elif self.color_space == 'rgba':
                rgb = rgba[:, :, :3]
                if self.is_show_normalization:
                    rgb = rgb.astype(np.float32) / 255.0
                src_img = rgb
            else:
                rgbf = rgba[:, :, :3].astype(np.float32) / 255.0
                hsv = cv2.cvtColor(rgbf, cv2.COLOR_RGB2HSV)
                src_img = hsv
            psnrs = image_metric.compute_psnr(src_img)
            psnrs_str = ', '.join([f"{psnr:.2f}" for psnr in psnrs])
            self.position_color_label.setText(f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), psnr(dB): ({psnrs_str})")

    def saveCorp(self, rect: QRectF, saveDir):
        qimage = self.image.copy(rect.toRect()).convertToFormat(QImage.Format_ARGB32)
        if not qimage.isNull():
            ptr = qimage.bits()
            ptr.setsize(qimage.byteCount())
            arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
            cv2.imwrite(os.path.join(saveDir, os.path.basename(self.image_path)), arr)

    def showPixelNormalization(self):
        self.is_show_normalization = not self.is_show_normalization
        if self.select_pos is not None:
            self.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            self.updateRectangle(self.select_rect, statistics=True)

    def updateColorSpace(self, color_space):
        self.color_space = color_space
        if self.select_pos is not None:
            self.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            self.updateRectangle(self.select_rect, statistics=True)


class ImageComparator(QMainWindow):
    def __init__(self, image_paths):
        super().__init__()
        image_paths.sort()
        self.image_paths = image_paths[:9]
        self.color_space = 'yuv'
        self.viewers = []
        self.grid_layout = None
        self.container = None
        self.initUI()
        self.select_pos = None
        self.select_rect: QRectF = None

        # 启用拖拽接收
        self.setAcceptDrops(True)

    def get_grid_layout(self, count):
        if count == 1:
            return (1, 1)
        elif count == 2:
            return (1, 2)
        elif count == 3:
            return (1, 3)
        elif count == 4:
            return (2, 2)
        elif count in (5, 6):
            return (3, 2)
        elif count in (7, 8, 9):
            return (3, 3)
        else:
            return (1, 1)

    def initUI(self):
        self.setWindowTitle('Image Comparator')
        self.resize(1200, 800)
        main_layout = QVBoxLayout()

        self.grid_layout = QGridLayout()
        self.viewers = []

        rows, cols = self.get_grid_layout(len(self.image_paths))
        for i, image_path in enumerate(self.image_paths):
            viewer = ImageViewer(image_path, os.path.basename(image_path))
            self.viewers.append(viewer)
            self.grid_layout.addWidget(viewer, i // cols, i % cols)
            viewer.offset_scale_changed.connect(self.offset_scale_sync)
            viewer.position_changed.connect(self.position_sync)
            viewer.rectangle_changed.connect(self.rectangle_sync)

        self.container = QWidget()
        self.container.setLayout(self.grid_layout)
        main_layout.addWidget(self.container)
        main_widget = QWidget()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        menubar = self.menuBar()
        file_menu = menubar.addMenu("文件(&F)")
        act_add = file_menu.addAction("添加图片(&A)...")
        act_add.triggered.connect(self.menu_add_images)
        act_clear = file_menu.addAction("清空图片(&C)")
        act_clear.triggered.connect(self.clearImages)
        act_export = file_menu.addAction("导出裁剪图(&E)...")
        act_export.triggered.connect(self.export_corp)
        act_reload = file_menu.addAction("重新加载(&R)")
        act_reload.setShortcut("F5")
        act_reload.triggered.connect(self.reloadImages)

        help_menu = menubar.addMenu("帮助(&H)")
        act_usage = help_menu.addAction("使用说明(&U)...")
        act_usage.triggered.connect(self.show_usage_dialog)

        self._toast = ToastWidget(self)

    def show_usage_dialog(self):
        prog_name = os.path.basename(sys.argv[0]) if sys.argv and sys.argv[0] else "image_comparator.py"
        text = usage(prog_name)
        dlg = UsageDialog(text, self)
        dlg.exec_()

    def show_toast(self, message: str, ms: int = 1000):
        if hasattr(self, "_toast") and self._toast is not None:
            self._toast.show_message(message, ms)

    def offset_scale_sync(self, offset: QPointF, scale: float):
        for viewer in self.viewers:
            viewer.updateGeometry(offset, scale)

    def position_sync(self, pos: QPointF):
        self.select_rect = None
        self.select_pos = pos
        for viewer in self.viewers:
            viewer.updatePosition(pos)

    def rectangle_sync(self, rect: QRectF):
        self.select_rect = rect
        for viewer in self.viewers:
            viewer.updateRectangle(rect, statistics=True)

    def calculate_hist(self):
        self.show_toast("计算直方图")
        if self.select_rect is not None:
            for viewer in self.viewers:
                viewer.updateHist(self.select_rect)

    def calculate_mtf(self):
        self.show_toast("计算 MTF")
        for viewer in self.viewers:
            viewer.updateMTF(self.select_rect)

    def calculate_psnr(self):
        self.show_toast("计算 峰值信噪比")
        for viewer in self.viewers:
            viewer.updatePSNR(self.select_rect)

    def export_corp(self):
        self.show_toast("导出剪切图")
        if not self.viewers:
            return
        saveDir = QFileDialog.getExistingDirectory(self.viewers[0], "Select Export Corp Directory")
        if self.select_rect is not None and saveDir:
            for viewer in self.viewers:
                viewer.saveCorp(self.select_rect, saveDir)

    def pixel_normalization(self):
        self.show_toast("切换像素归一化\原始数据显示")
        for viewer in self.viewers:
            viewer.showPixelNormalization()

    def toggleColorSpace(self):
        if self.color_space == 'yuv':
            self.color_space = 'rgba'
        elif self.color_space == 'rgba':
            self.color_space = 'hsv'
        else:
            self.color_space = 'yuv'
        self.show_toast(f"切换到{self.color_space}颜色空间")
        for viewer in self.viewers:
            viewer.updateColorSpace(self.color_space)

    def rotateLeft(self):
        self.show_toast("逆时针旋转90度")
        for viewer in self.viewers:
            viewer.rotateImage(clockwise=False)
        self.select_pos = None
        self.select_rect = None

    def rotateRight(self):
        self.show_toast("顺时针旋转90度")
        for viewer in self.viewers:
            viewer.rotateImage(clockwise=True)
        self.select_pos = None
        self.select_rect = None

    def keyPressEvent(self, event: QKeyEvent):
        key = event.key()
        if key == Qt.Key_M:
            self.calculate_mtf()
        elif key == Qt.Key_H:
            self.calculate_hist()
        elif key == Qt.Key_E:
            self.export_corp()
        elif key == Qt.Key_N:
            self.pixel_normalization()
        elif key == Qt.Key_C:
            self.toggleColorSpace()
        elif key == Qt.Key_L:
            self.rotateLeft()
        elif key == Qt.Key_R:
            self.rotateRight()
        elif key == Qt.Key_P:
            self.calculate_psnr()
        elif key == Qt.Key_F5:
            self.reloadImages()

    def _reflow_grid(self):
        while self.grid_layout.count():
            item = self.grid_layout.takeAt(0)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)

        rows, cols = self.get_grid_layout(len(self.viewers))
        for i, viewer in enumerate(self.viewers):
            self.grid_layout.addWidget(viewer, i // cols, i % cols)

        self.container.setLayout(self.grid_layout)
        self.container.update()
        self.container.repaint()

    def addImage(self, image_paths):
        """
        动态添加图片（1~N），自动更新网格布局与信号。
        - 忽略无效路径与重复路径
        - 总数最多 9 张
        - 添加前对 all_paths 排序
        """
        if not image_paths:
            return

        new_paths = [p for p in image_paths if os.path.isfile(p)]
        if not new_paths:
            return

        existed = set(self.image_paths)
        new_paths = [p for p in new_paths if p not in existed]
        if not new_paths:
            return

        all_paths = self.image_paths + new_paths
        # 排序
        all_paths = sorted(all_paths)
        if len(all_paths) > 9:
            all_paths = all_paths[:9]

        keep_set = set(all_paths)
        # 重新计算 new_paths，以排序和截断后保持一致
        new_paths = [p for p in all_paths if p not in set(self.image_paths)]

        self.image_paths = all_paths

        # 清空并按排序后的路径重建 viewers，确保排序与布局一致
        # 保留已有 viewer 对应路径，避免重复创建
        old_by_path = {v.image_path: v for v in self.viewers}
        new_viewers = []
        for p in self.image_paths:
            if p in old_by_path:
                new_viewers.append(old_by_path[p])
            else:
                viewer = ImageViewer(p, os.path.basename(p))
                viewer.updateColorSpace(self.color_space)
                viewer.offset_scale_changed.connect(self.offset_scale_sync)
                viewer.position_changed.connect(self.position_sync)
                viewer.rectangle_changed.connect(self.rectangle_sync)
                new_viewers.append(viewer)
        self.viewers = new_viewers

        self._reflow_grid()

        # 同步当前选择状态
        if self.select_pos is not None:
            for v in self.viewers:
                v.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            for v in self.viewers:
                v.updateRectangle(self.select_rect, statistics=True)

    def reloadImages(self):
        """
        重新加载当前 image_paths 中的图片：
        - 清空现有 viewers 与布局
        - 依据 image_paths 重建 viewers
        - 重新布局
        - 恢复 color_space
        - 清空选择状态（符合“清理窗口后重新load”的要求）
        """
        # 1) 清空布局与旧视图
        self.show_toast("重新加载图片")
        for i in reversed(range(self.grid_layout.count())):
            item = self.grid_layout.takeAt(i)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)
                w.setParent(None)
        self.viewers.clear()

        # 2) 重建 viewers
        rows, cols = self.get_grid_layout(len(self.image_paths))
        for i, image_path in enumerate(self.image_paths):
            viewer = ImageViewer(image_path, os.path.basename(image_path))
            # 恢复颜色空间
            viewer.updateColorSpace(self.color_space)
            # 连接信号
            viewer.offset_scale_changed.connect(self.offset_scale_sync)
            viewer.position_changed.connect(self.position_sync)
            viewer.rectangle_changed.connect(self.rectangle_sync)
            self.viewers.append(viewer)
            self.grid_layout.addWidget(viewer, i // cols, i % cols)

        self.container.setLayout(self.grid_layout)
        self.container.update()
        self.container.repaint()

        # 3) 清空选择状态（符合“清理窗口后重新load”）
        self.select_pos = None
        self.select_rect = None


    def clearImages(self):
        """移除所有图片与视图"""
        for i in reversed(range(self.grid_layout.count())):
            item = self.grid_layout.takeAt(i)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)
                w.setParent(None)
        self.viewers.clear()
        self.image_paths = []
        self.select_pos = None
        self.select_rect = None
        self._reflow_grid()

    def menu_add_images(self):
        paths, _ = QFileDialog.getOpenFileNames(
            self,
            "选择要添加的图片",
            "",
            "Images (*.png *.jpg *.jpeg *.bmp *.tif *.tiff *.webp);;All Files (*)"
        )
        if paths:
            self.addImage(paths)

    # 拖拽文件到窗口添加
    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            # 只要有本地文件就接受
            urls = event.mimeData().urls()
            local_files = [u.toLocalFile() for u in urls if u.isLocalFile()]
            if local_files:
                event.acceptProposedAction()
                return
        event.ignore()

    def dropEvent(self, event: QDropEvent):
        urls = event.mimeData().urls()
        local_files = [u.toLocalFile() for u in urls if u.isLocalFile()]
        if local_files:
            self.addImage(local_files)
            event.acceptProposedAction()
        else:
            event.ignore()


def usage(prog_name: str = "image_comparator.py") -> str:
    return f"""
用法:
python {prog_name} <image1> <image2> ... <imageN>

描述:
并排对比最多 9 张图片，支持同步缩放/拖拽、点与矩形选区、颜色信息查看、
直方图统计、MTF 计算、裁剪导出、像素归一化显示、颜色空间切换以及旋转。

示例:
python {prog_name} a.png b.jpg c.bmp

快捷键:
M  - 计算并显示选区的 MTF 信息
H  - 计算并显示选区的直方图
E  - 将当前选区导出另存为crop图片（弹出目录选择）
N  - 切换像素是否归一化显示(0-1)
C  - 切换颜色空间（YUV <-> RGBA <-> HSV）
P  - 计算并显示选区的 PSNR 信息
S  - 计算并显示选区的饱和度信息（预留）
L  - 向左旋转 90 度
R  - 向右旋转 90 度

鼠标操作:
左键点击：点选，显示该点颜色信息
左键拖拽：框选矩形区域
右键拖拽：平移
滚轮：以视图中心为基准缩放

参数:
images (位置参数)  输入的图片路径，1~9 张。

注意:
- 当前最多支持 9 张图片。
"""


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        add_help=False,
        prog="image_comparator.py"
    )
    parser.add_argument('-h', '--help', action='store_true', help='显示帮助信息并退出')
    parser.add_argument('images', nargs='*', help='图片路径，最多 9 张')

    args = parser.parse_args()

    if args.help:
        print(usage(parser.prog))
        sys.exit(0)

    image_paths = args.images[:9] if args.images else []

    invalid = [p for p in image_paths if not os.path.isfile(p)]
    if invalid:
        print("下列路径不是有效文件：")
        for p in invalid:
            print(f"  - {p}")
        print()
        print(usage(parser.prog))
        sys.exit(1)

    app = QApplication(sys.argv)
    comparator = ImageComparator(image_paths)
    comparator.show()
    sys.exit(app.exec_())
